import json
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad


def vn_function(key_str):
    """
    将输入字符串转换为16字节的密钥
    相当于JavaScript中的Vn函数
    """
    # 将字符串转换为UTF-8字节
    utf8_bytes = key_str.encode('utf-8')

    # 创建16字节的数组，不足的用0填充
    result = bytearray(16)
    for i in range(min(len(utf8_bytes), 16)):
        result[i] = utf8_bytes[i]

    return bytes(result)


def encrypt_data(data, key_str="fsdfjk12390jqIOJSDKL"):
    """
    加密函数，与decrypt_data相对应

    参数:
        data: 要加密的数据（JSON对象或字典）
        key_str: AES密钥

    返回:
        加密后的字符串（URL安全的Base64格式）
    """
    # 如果输入为空，直接返回
    if data is None:
        return None

    try:
        # 将数据转换为JSON字符串
        json_str = json.dumps(data, ensure_ascii=False)

        # 将字符串转换为UTF-8字节
        json_bytes = json_str.encode('utf-8')

        # 处理密钥
        key = vn_function(key_str)

        # 创建AES-ECB加密器
        cipher = AES.new(key, AES.MODE_ECB)

        # 加密并添加填充
        encrypted_bytes = cipher.encrypt(pad(json_bytes, AES.block_size))

        # Base64编码
        standard_base64 = base64.b64encode(encrypted_bytes).decode('utf-8')

        # 将标准Base64转换为URL安全的Base64
        url_safe_base64 = standard_base64.replace('+', '-').replace('/', '_')

        return url_safe_base64
    except Exception as e:
        print(f"加密失败: {e}")
        return None


if __name__ == "__main__":
    # 要加密的数据文件路径
    input_path = "./output/decrypted_data.json"

    # 加密数据导出地址
    output_path = "./input/encrypted_data_new.txt"

    # 读取JSON数据
    with open(input_path, 'r', encoding='utf-8') as file:
        data = json.load(file)

    # 加密数据
    encrypted_data = encrypt_data(data)

    # 将加密后的数据保存到文件
    with open(output_path, 'w', encoding='utf-8') as file:
        file.write(encrypted_data)

    print("数据已成功加密并保存到:", output_path)
    print("加密后的数据:", encrypted_data)
